{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: torch==2.2 in /opt/conda/lib/python3.10/site-packages (2.2.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (4.9.0)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (3.1.2)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2023.12.2)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2.19.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (12.1.105)\n",
      "Requirement already satisfied: triton==2.2.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2) (2.2.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /opt/conda/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch==2.2) (12.3.101)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch==2.2) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch==2.2) (1.3.0)\n",
      "Requirement already satisfied: torchvision==0.17 in /opt/conda/lib/python3.10/site-packages (0.17.0)\n",
      "Requirement already satisfied: numpy in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17) (1.26.2)\n",
      "Requirement already satisfied: requests in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17) (2.28.1)\n",
      "Requirement already satisfied: torch==2.2.0 in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17) (2.2.0)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /opt/conda/lib/python3.10/site-packages (from torchvision==0.17) (10.2.0)\n",
      "Requirement already satisfied: filelock in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (4.9.0)\n",
      "Requirement already satisfied: sympy in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (1.12)\n",
      "Requirement already satisfied: networkx in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (3.1.2)\n",
      "Requirement already satisfied: fsspec in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (2023.12.2)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.19.3 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (2.19.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (12.1.105)\n",
      "Requirement already satisfied: triton==2.2.0 in /opt/conda/lib/python3.10/site-packages (from torch==2.2.0->torchvision==0.17) (2.2.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /opt/conda/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch==2.2.0->torchvision==0.17) (12.3.101)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17) (2.1.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17) (3.3)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17) (1.26.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests->torchvision==0.17) (2022.6.15)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->torch==2.2.0->torchvision==0.17) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in /opt/conda/lib/python3.10/site-packages (from sympy->torch==2.2.0->torchvision==0.17) (1.3.0)\n",
      "Requirement already satisfied: matplotlib==3.5.2 in /opt/conda/lib/python3.10/site-packages (3.5.2)\n",
      "Requirement already satisfied: cycler>=0.10 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (4.46.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (1.4.5)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (1.26.2)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (21.3)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (10.2.0)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (3.0.9)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /opt/conda/lib/python3.10/site-packages (from matplotlib==3.5.2) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib==3.5.2) (1.16.0)\n",
      "Requirement already satisfied: scikit-image==0.19.3 in /opt/conda/lib/python3.10/site-packages (0.19.3)\n",
      "Requirement already satisfied: numpy>=1.17.0 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (1.26.2)\n",
      "Requirement already satisfied: scipy>=1.4.1 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (1.11.4)\n",
      "Requirement already satisfied: networkx>=2.2 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (2.8.5)\n",
      "Requirement already satisfied: pillow!=7.1.0,!=7.1.1,!=8.3.0,>=6.1.0 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (10.2.0)\n",
      "Requirement already satisfied: imageio>=2.4.1 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (2.33.0)\n",
      "Requirement already satisfied: tifffile>=2019.7.26 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (2023.12.9)\n",
      "Requirement already satisfied: PyWavelets>=1.1.1 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (1.5.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from scikit-image==0.19.3) (21.3)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->scikit-image==0.19.3) (3.0.9)\n"
     ]
    }
   ],
   "source": [
    "!pip install torch==2.2\n",
    "!pip install torchvision==0.17\n",
    "!pip install matplotlib==3.5.2\n",
    "!pip install scikit-image==0.19.3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import required librarires"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.autograd import Variable\n",
    "\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.utils import save_image\n",
    "from torchvision import datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define constants / model hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_eps=10\n",
    "bsize=32\n",
    "lrate=0.001\n",
    "lat_dimension=64\n",
    "image_sz=64\n",
    "chnls=1\n",
    "logging_intv=200"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GANGenerator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GANGenerator, self).__init__()\n",
    "        self.inp_sz = image_sz // 4\n",
    "        self.lin = nn.Linear(lat_dimension, 128 * self.inp_sz ** 2)\n",
    "        self.bn1 = nn.BatchNorm2d(128)\n",
    "        self.up1 = nn.Upsample(scale_factor=2)\n",
    "        self.cn1 = nn.Conv2d(128, 128, 3, stride=1, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(128, 0.8)\n",
    "        self.rl1 = nn.LeakyReLU(0.2, inplace=True)\n",
    "        self.up2 = nn.Upsample(scale_factor=2)\n",
    "        self.cn2 = nn.Conv2d(128, 64, 3, stride=1, padding=1)\n",
    "        self.bn3 = nn.BatchNorm2d(64, 0.8)\n",
    "        self.rl2 = nn.LeakyReLU(0.2, inplace=True)\n",
    "        self.cn3 = nn.Conv2d(64, chnls, 3, stride=1, padding=1)\n",
    "        self.act = nn.Tanh()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.lin(x)\n",
    "        x = x.view(x.shape[0], 128, self.inp_sz, self.inp_sz)\n",
    "        x = self.bn1(x)\n",
    "        x = self.up1(x)\n",
    "        x = self.cn1(x)\n",
    "        x = self.bn2(x)\n",
    "        x = self.rl1(x)\n",
    "        x = self.up2(x)\n",
    "        x = self.cn2(x)\n",
    "        x = self.bn3(x)\n",
    "        x = self.rl2(x)\n",
    "        x = self.cn3(x)\n",
    "        out = self.act(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GANDiscriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GANDiscriminator, self).__init__()\n",
    "\n",
    "        def disc_module(ip_chnls, op_chnls, bnorm=True):\n",
    "            mod = [nn.Conv2d(ip_chnls, op_chnls, 3, 2, 1), \n",
    "                   nn.LeakyReLU(0.2, inplace=True), \n",
    "                   nn.Dropout2d(0.25)]\n",
    "            if bnorm:\n",
    "                mod += [nn.BatchNorm2d(op_chnls, 0.8)]\n",
    "            return mod\n",
    "\n",
    "        self.disc_model = nn.Sequential(\n",
    "            *disc_module(chnls, 16, bnorm=False),\n",
    "            *disc_module(16, 32),\n",
    "            *disc_module(32, 64),\n",
    "            *disc_module(64, 128),\n",
    "        )\n",
    "\n",
    "        # width and height of the down-sized image\n",
    "        ds_size = image_sz // 2 ** 4\n",
    "        self.adverse_lyr = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.disc_model(x)\n",
    "        x = x.view(x.shape[0], -1)\n",
    "        out = self.adverse_lyr(x)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# instantiate the discriminator and generator models\n",
    "gen = GANGenerator()\n",
    "disc = GANDiscriminator()\n",
    "\n",
    "# define the loss metric\n",
    "adv_loss_func = torch.nn.BCELoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the image dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./data/mnist/MNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9912422/9912422 [00:00<00:00, 205995755.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./data/mnist/MNIST/raw/train-images-idx3-ubyte.gz to ./data/mnist/MNIST/raw\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./data/mnist/MNIST/raw/train-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 28881/28881 [00:00<00:00, 124241737.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./data/mnist/MNIST/raw/train-labels-idx1-ubyte.gz to ./data/mnist/MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./data/mnist/MNIST/raw/t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1648877/1648877 [00:00<00:00, 66561036.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./data/mnist/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data/mnist/MNIST/raw\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./data/mnist/MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4542/4542 [00:00<00:00, 24644927.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./data/mnist/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/mnist/MNIST/raw\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# define the dataset and corresponding dataloader\n",
    "dloader = torch.utils.data.DataLoader(\n",
    "    datasets.MNIST(\n",
    "        \"./data/mnist/\",\n",
    "        download=True,\n",
    "        transform=transforms.Compose(\n",
    "            [transforms.Resize((image_sz, image_sz)), \n",
    "             transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]\n",
    "        ),\n",
    "    ),\n",
    "    batch_size=bsize,\n",
    "    shuffle=True,\n",
    ")\n",
    "\n",
    "# define the optimization schedule for both G and D\n",
    "opt_gen = torch.optim.Adam(gen.parameters(), lr=lrate)\n",
    "opt_disc = torch.optim.Adam(disc.parameters(), lr=lrate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training loop for DCGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch number 0 | batch number 0 | generator loss = 0.6827291250228882 | discriminator loss = 0.6916588544845581\n",
      "epoch number 0 | batch number 200 | generator loss = 1.3753488063812256 | discriminator loss = 0.5124980211257935\n",
      "epoch number 0 | batch number 400 | generator loss = 0.8567953705787659 | discriminator loss = 0.4100569188594818\n",
      "epoch number 0 | batch number 600 | generator loss = 1.2081809043884277 | discriminator loss = 0.3683425486087799\n",
      "epoch number 0 | batch number 800 | generator loss = 1.4627505540847778 | discriminator loss = 0.5539621114730835\n",
      "epoch number 0 | batch number 1000 | generator loss = 2.522969961166382 | discriminator loss = 0.22689306735992432\n",
      "epoch number 0 | batch number 1200 | generator loss = 0.7817821502685547 | discriminator loss = 0.47197943925857544\n",
      "epoch number 0 | batch number 1400 | generator loss = 1.5259451866149902 | discriminator loss = 0.7356035709381104\n",
      "epoch number 0 | batch number 1600 | generator loss = 2.405836820602417 | discriminator loss = 0.3994973301887512\n",
      "epoch number 0 | batch number 1800 | generator loss = 1.8351649045944214 | discriminator loss = 0.48325562477111816\n",
      "epoch number 1 | batch number 125 | generator loss = 3.6158790588378906 | discriminator loss = 0.532764732837677\n",
      "epoch number 1 | batch number 325 | generator loss = 5.781711101531982 | discriminator loss = 0.4379121959209442\n",
      "epoch number 1 | batch number 525 | generator loss = 1.9844849109649658 | discriminator loss = 0.44709718227386475\n",
      "epoch number 1 | batch number 725 | generator loss = 4.3647541999816895 | discriminator loss = 0.8382872343063354\n",
      "epoch number 1 | batch number 925 | generator loss = 6.970343112945557 | discriminator loss = 0.27220046520233154\n",
      "epoch number 1 | batch number 1125 | generator loss = 2.351853370666504 | discriminator loss = 0.13503685593605042\n",
      "epoch number 1 | batch number 1325 | generator loss = 0.5296928286552429 | discriminator loss = 0.4074157178401947\n",
      "epoch number 1 | batch number 1525 | generator loss = 2.7166566848754883 | discriminator loss = 0.1553489714860916\n",
      "epoch number 1 | batch number 1725 | generator loss = 4.88264274597168 | discriminator loss = 0.06121131032705307\n",
      "epoch number 2 | batch number 50 | generator loss = 2.3263041973114014 | discriminator loss = 0.121161088347435\n",
      "epoch number 2 | batch number 250 | generator loss = 1.4021943807601929 | discriminator loss = 0.1081303060054779\n",
      "epoch number 2 | batch number 450 | generator loss = 4.399090766906738 | discriminator loss = 0.2777940332889557\n",
      "epoch number 2 | batch number 650 | generator loss = 4.517523765563965 | discriminator loss = 0.16361413896083832\n",
      "epoch number 2 | batch number 850 | generator loss = 3.468811511993408 | discriminator loss = 0.6588649749755859\n",
      "epoch number 2 | batch number 1050 | generator loss = 4.116889953613281 | discriminator loss = 0.3950495421886444\n",
      "epoch number 2 | batch number 1250 | generator loss = 3.0756163597106934 | discriminator loss = 0.3150275647640228\n",
      "epoch number 2 | batch number 1450 | generator loss = 1.878197193145752 | discriminator loss = 0.1765824407339096\n",
      "epoch number 2 | batch number 1650 | generator loss = 4.071030139923096 | discriminator loss = 0.3719501197338104\n",
      "epoch number 2 | batch number 1850 | generator loss = 4.845463275909424 | discriminator loss = 0.5089448690414429\n",
      "epoch number 3 | batch number 175 | generator loss = 1.359919786453247 | discriminator loss = 0.6091972589492798\n",
      "epoch number 3 | batch number 375 | generator loss = 5.107836723327637 | discriminator loss = 0.03625940531492233\n",
      "epoch number 3 | batch number 575 | generator loss = 1.296502947807312 | discriminator loss = 0.23591016232967377\n",
      "epoch number 3 | batch number 775 | generator loss = 2.5184953212738037 | discriminator loss = 0.2834915220737457\n",
      "epoch number 3 | batch number 975 | generator loss = 3.4772181510925293 | discriminator loss = 0.119707852602005\n",
      "epoch number 3 | batch number 1175 | generator loss = 3.1785566806793213 | discriminator loss = 0.09721270948648453\n",
      "epoch number 3 | batch number 1375 | generator loss = 0.8171387910842896 | discriminator loss = 0.5593546628952026\n",
      "epoch number 3 | batch number 1575 | generator loss = 4.523009300231934 | discriminator loss = 0.20698320865631104\n",
      "epoch number 3 | batch number 1775 | generator loss = 5.275208950042725 | discriminator loss = 0.029274985194206238\n",
      "epoch number 4 | batch number 100 | generator loss = 4.304698944091797 | discriminator loss = 0.13636040687561035\n",
      "epoch number 4 | batch number 300 | generator loss = 2.9817628860473633 | discriminator loss = 0.43162867426872253\n",
      "epoch number 4 | batch number 500 | generator loss = 7.312787055969238 | discriminator loss = 0.03579274192452431\n",
      "epoch number 4 | batch number 700 | generator loss = 2.0278713703155518 | discriminator loss = 0.10602187365293503\n",
      "epoch number 4 | batch number 900 | generator loss = 3.584226131439209 | discriminator loss = 0.6118831038475037\n",
      "epoch number 4 | batch number 1100 | generator loss = 6.313017845153809 | discriminator loss = 0.00797022320330143\n",
      "epoch number 4 | batch number 1300 | generator loss = 6.951356887817383 | discriminator loss = 0.04062850773334503\n",
      "epoch number 4 | batch number 1500 | generator loss = 7.487911224365234 | discriminator loss = 0.10682754963636398\n",
      "epoch number 4 | batch number 1700 | generator loss = 3.9839158058166504 | discriminator loss = 0.1216537281870842\n",
      "epoch number 5 | batch number 25 | generator loss = 6.013751029968262 | discriminator loss = 0.14520937204360962\n",
      "epoch number 5 | batch number 225 | generator loss = 3.639707088470459 | discriminator loss = 0.26481893658638\n",
      "epoch number 5 | batch number 425 | generator loss = 2.0560028553009033 | discriminator loss = 0.47513139247894287\n",
      "epoch number 5 | batch number 625 | generator loss = 2.8454513549804688 | discriminator loss = 0.03515440225601196\n",
      "epoch number 5 | batch number 825 | generator loss = 1.5685524940490723 | discriminator loss = 0.06366463750600815\n",
      "epoch number 5 | batch number 1025 | generator loss = 6.179746627807617 | discriminator loss = 0.005660687107592821\n",
      "epoch number 5 | batch number 1225 | generator loss = 4.440589904785156 | discriminator loss = 0.20085813105106354\n",
      "epoch number 5 | batch number 1425 | generator loss = 5.666544437408447 | discriminator loss = 0.11899282783269882\n",
      "epoch number 5 | batch number 1625 | generator loss = 2.9011645317077637 | discriminator loss = 0.02219812013208866\n",
      "epoch number 5 | batch number 1825 | generator loss = 3.602222442626953 | discriminator loss = 0.17869994044303894\n",
      "epoch number 6 | batch number 150 | generator loss = 2.7667925357818604 | discriminator loss = 0.09315657615661621\n",
      "epoch number 6 | batch number 350 | generator loss = 3.399181604385376 | discriminator loss = 0.05052967369556427\n",
      "epoch number 6 | batch number 550 | generator loss = 3.652853488922119 | discriminator loss = 0.060660507529973984\n",
      "epoch number 6 | batch number 750 | generator loss = 8.414639472961426 | discriminator loss = 0.018418243154883385\n",
      "epoch number 6 | batch number 950 | generator loss = 5.858670711517334 | discriminator loss = 0.1595914363861084\n",
      "epoch number 6 | batch number 1150 | generator loss = 3.801039218902588 | discriminator loss = 0.24717527627944946\n",
      "epoch number 6 | batch number 1350 | generator loss = 1.436686396598816 | discriminator loss = 0.27819910645484924\n",
      "epoch number 6 | batch number 1550 | generator loss = 5.327937602996826 | discriminator loss = 0.038817644119262695\n",
      "epoch number 6 | batch number 1750 | generator loss = 3.9502553939819336 | discriminator loss = 0.21309791505336761\n",
      "epoch number 7 | batch number 75 | generator loss = 6.7195539474487305 | discriminator loss = 0.1338636875152588\n",
      "epoch number 7 | batch number 275 | generator loss = 2.8770134449005127 | discriminator loss = 0.08344466984272003\n",
      "epoch number 7 | batch number 475 | generator loss = 5.707181930541992 | discriminator loss = 0.24805203080177307\n",
      "epoch number 7 | batch number 675 | generator loss = 4.36092472076416 | discriminator loss = 0.013633709400892258\n",
      "epoch number 7 | batch number 875 | generator loss = 7.694408416748047 | discriminator loss = 0.061054777354002\n",
      "epoch number 7 | batch number 1075 | generator loss = 2.6442079544067383 | discriminator loss = 0.03422663360834122\n",
      "epoch number 7 | batch number 1275 | generator loss = 6.279775619506836 | discriminator loss = 0.0639989823102951\n",
      "epoch number 7 | batch number 1475 | generator loss = 2.6854538917541504 | discriminator loss = 0.015679141506552696\n",
      "epoch number 7 | batch number 1675 | generator loss = 6.811121940612793 | discriminator loss = 0.16857674717903137\n",
      "epoch number 8 | batch number 0 | generator loss = 2.719939708709717 | discriminator loss = 0.13963723182678223\n",
      "epoch number 8 | batch number 200 | generator loss = 3.630988597869873 | discriminator loss = 0.08019018918275833\n",
      "epoch number 8 | batch number 400 | generator loss = 2.3935306072235107 | discriminator loss = 0.07276156544685364\n",
      "epoch number 8 | batch number 600 | generator loss = 7.66723108291626 | discriminator loss = 0.026261799037456512\n",
      "epoch number 8 | batch number 800 | generator loss = 5.463866233825684 | discriminator loss = 0.0509343259036541\n",
      "epoch number 8 | batch number 1000 | generator loss = 7.421459197998047 | discriminator loss = 0.3396588861942291\n",
      "epoch number 8 | batch number 1200 | generator loss = 4.4102091789245605 | discriminator loss = 0.35801374912261963\n",
      "epoch number 8 | batch number 1400 | generator loss = 7.551414489746094 | discriminator loss = 0.5030229091644287\n",
      "epoch number 8 | batch number 1600 | generator loss = 7.160701274871826 | discriminator loss = 0.25210389494895935\n",
      "epoch number 8 | batch number 1800 | generator loss = 8.119586944580078 | discriminator loss = 0.07027419656515121\n",
      "epoch number 9 | batch number 125 | generator loss = 3.109663963317871 | discriminator loss = 0.012378710322082043\n",
      "epoch number 9 | batch number 325 | generator loss = 3.354196310043335 | discriminator loss = 0.2538221478462219\n",
      "epoch number 9 | batch number 525 | generator loss = 2.339954137802124 | discriminator loss = 0.2265104055404663\n",
      "epoch number 9 | batch number 725 | generator loss = 5.504094123840332 | discriminator loss = 0.09025933593511581\n",
      "epoch number 9 | batch number 925 | generator loss = 2.890489101409912 | discriminator loss = 0.8024916648864746\n",
      "epoch number 9 | batch number 1125 | generator loss = 6.0101423263549805 | discriminator loss = 0.01851765438914299\n",
      "epoch number 9 | batch number 1325 | generator loss = 8.085104942321777 | discriminator loss = 0.2512941062450409\n",
      "epoch number 9 | batch number 1525 | generator loss = 2.770570755004883 | discriminator loss = 0.11811449378728867\n",
      "epoch number 9 | batch number 1725 | generator loss = 7.662922382354736 | discriminator loss = 0.0023379907943308353\n"
     ]
    }
   ],
   "source": [
    "os.makedirs(\"./images_mnist\", exist_ok=True)\n",
    "\n",
    "for ep in range(num_eps):\n",
    "    for idx, (images, _) in enumerate(dloader):\n",
    "\n",
    "        # generate grounnd truths for real and fake images\n",
    "        good_img = Variable(torch.FloatTensor(images.shape[0], 1).fill_(1.0), requires_grad=False)\n",
    "        bad_img = Variable(torch.FloatTensor(images.shape[0], 1).fill_(0.0), requires_grad=False)\n",
    "\n",
    "        # get a real image\n",
    "        actual_images = Variable(images.type(torch.FloatTensor))\n",
    "\n",
    "        # train the generator model\n",
    "        opt_gen.zero_grad()\n",
    "\n",
    "        # generate a batch of images based on random noise as input\n",
    "        noise = Variable(torch.FloatTensor(np.random.normal(0, 1, (images.shape[0], lat_dimension))))\n",
    "        gen_images = gen(noise)\n",
    "\n",
    "        # generator model optimization - how well can it fool the discriminator\n",
    "        generator_loss = adv_loss_func(disc(gen_images), good_img)\n",
    "        generator_loss.backward()\n",
    "        opt_gen.step()\n",
    "\n",
    "        # train the discriminator model\n",
    "        opt_disc.zero_grad()\n",
    "\n",
    "        # calculate discriminator loss as average of mistakes(losses) in confusing real images as fake and vice versa\n",
    "        actual_image_loss = adv_loss_func(disc(actual_images), good_img)\n",
    "        fake_image_loss = adv_loss_func(disc(gen_images.detach()), bad_img)\n",
    "        discriminator_loss = (actual_image_loss + fake_image_loss) / 2\n",
    "\n",
    "        # discriminator model optimization\n",
    "        discriminator_loss.backward()\n",
    "        opt_disc.step()\n",
    "\n",
    "        batches_completed = ep * len(dloader) + idx\n",
    "        if batches_completed % logging_intv == 0:\n",
    "            print(f\"epoch number {ep} | batch number {idx} | generator loss = {generator_loss.item()} | discriminator loss = {discriminator_loss.item()}\")\n",
    "            save_image(gen_images.data[:25], f\"images_mnist/{batches_completed}.png\", nrow=5, normalize=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (Local)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
